{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "W12_Tutorial1",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/CIS-522/course-content/blob/w12_t1_prepod/tutorials/W12_DeepRL/W12_Tutorial1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2g-b2Yq3N2nK"
      },
      "source": [
        "# CIS-522 Week 12 Part 1\n",
        "# Policy Gradient and Model-Based RL\n",
        "\n",
        "__Instructor:__ Dinesh Jayaraman\n",
        "\n",
        "__Content creators:__ Varun Velpula, Chuning Zhu\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Tg3o4tUiduU8"
      },
      "source": [
        "#@markdown What is your Pennkey and pod? (text, not numbers, e.g. bfranklin)\n",
        "my_pennkey = '' #@param {type:\"string\"}\n",
        "my_pod = 'Select' #@param ['Select', 'euclidean-wombat', 'sublime-newt', 'buoyant-unicorn', 'lackadaisical-manatee','indelible-stingray','superfluous-lyrebird','discreet-reindeer','quizzical-goldfish','astute-jellyfish','ubiquitous-cheetah','nonchalant-crocodile','fashionable-lemur','spiffy-eagle','electric-emu','quotidian-lion']\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0sDntL2ZPL3D"
      },
      "source": [
        "## Recap the experience from last week\n",
        "\n",
        "What did you learn last week. What questions do you have? [10 min discussion]"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "2Uuc48zoPKfv"
      },
      "source": [
        "learning_from_previous_week = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E7JpuimpPQoC"
      },
      "source": [
        "---\n",
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RFrct6E6P1ir",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Install\n",
        "\n",
        "!apt -q install imagemagick\n",
        "!apt install -q xvfb\n",
        "!pip install -q pyvirtualdisplay"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b9fHlT2agJdi"
      },
      "source": [
        "# imports\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import random\n",
        "import sys\n",
        "import time\n",
        "from builtins import super\n",
        "from collections import namedtuple\n",
        "from __future__ import print_function\n",
        "from pathlib import Path\n",
        "from tqdm.auto import tqdm\n",
        "\n",
        "\n",
        "# PyTorch\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable as V\n",
        "from torch.distributions import Normal\n",
        "\n",
        "# device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "device = torch.device(\"cpu\") # CPU actually trains faster because there is no batch\n",
        "\n",
        "# IPython\n",
        "import ipywidgets as widgets\n",
        "from ipywidgets import interact, interactive, fixed, interact_manual\n",
        "from IPython import display\n",
        "from IPython.display import Image # display\n",
        "from IPython.core.interactiveshell import InteractiveShell\n",
        "InteractiveShell.ast_node_interactivity = \"all\"\n",
        "\n",
        "# Gym\n",
        "import gym\n",
        "import gym.spaces\n",
        "gym.logger.set_level(gym.logger.ERROR)\n",
        "\n",
        "# Gym rendering\n",
        "from pyvirtualdisplay import Display\n",
        "\n",
        "d = Display()\n",
        "d.start();\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tSDUhtJQPwZ8",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Figure setting\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import animation\n",
        "\n",
        "%matplotlib inline \n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "\n",
        "plt.rcParams.update(plt.rcParamsDefault)\n",
        "plt.rc('animation', html='jshtml')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zmLB0D3HPkDI",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Seeds\n",
        "seed = 1\n",
        "random.seed(seed)\n",
        "torch.manual_seed(seed)\n",
        "torch.cuda.manual_seed_all(seed)\n",
        "torch.cuda.manual_seed(seed);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "qwMxXTBzIWkw"
      },
      "source": [
        "#@markdown ## Helper functions\n",
        "\n",
        "def visualize(env, policy, ep_steps):\n",
        "    # Simulate agent in environment\n",
        "    frames = []\n",
        "    state = env.reset()\n",
        "    for s in range(ep_steps):\n",
        "        action, _ = policy.get_action(state)\n",
        "        next_state, _, done, _ = env.step(action)\n",
        "\n",
        "        display.clear_output(wait=True)\n",
        "        frames.append(env.render('rgb_array'))\n",
        "        state = next_state\n",
        "        if done:\n",
        "            break\n",
        "    \n",
        "    # Generate video\n",
        "    plt.figure(figsize=(frames[0].shape[1] / 72.0, frames[0].shape[0] / 72.0), dpi=72)\n",
        "    patch = plt.imshow(frames[0])\n",
        "    plt.axis('off')\n",
        "    animate = lambda i: patch.set_data(frames[i])\n",
        "    anim = animation.FuncAnimation(plt.gcf(), animate, frames=range(len(frames)), interval=50)\n",
        "    plt.close() # avoid showing extra plots\n",
        "    return anim"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pX8lPj9RvXfS"
      },
      "source": [
        "---\n",
        "# Overview\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YDjL3Snj9jp7",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Policy Gradient Introduction\n",
        "\n",
        "import time\n",
        "try: t0;\n",
        "except NameError: t0=time.time()\n",
        "\n",
        "from IPython.display import YouTubeVideo\n",
        "\n",
        "video = YouTubeVideo(id=\"Yx_7nSg979Q\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "228Kg8ysBqQD"
      },
      "source": [
        "Last week, we learned our first deep reinforcement learning algorithm -- deep Q-learning. Deep Q-learning has its root in tabular Q-learning, where we estimate the value of each action at a state and take the action with the maximum value. As we have seen, deep Q-learning is powerful and can be leveraged to solve Atari games from high-dimensional observations. However, it has a few drawbacks: \n",
        "\n",
        "- The policy implied by Q-Learning is deterministic. This means that Q-Learning can’t learn stochastic policies, which can be useful in some environments. It also means we need to create our own exploration strategy since following the policy will not perform any exploration. We usually do this with $\\epsilon$-greedy exploration which can be quite inefficient.\n",
        "- There is no straightforward way to handle continuous actions in Q-Learning. We can go around this by having an additional actor network which predicts the argmax action. This leads to the DDPG algorithm.\n",
        "\n",
        "Recall that the goal of reinforcement learning is to learn an optimal **policy**. Can we do that directly? In this tutorial, we will introduce another family of RL algorithms called policy gradient methods, which optimize the policy directly using gradients of the RL objective. Here are a few reasons why policy gradient can be desirable:\n",
        "\n",
        "\n",
        "- In policy gradient, handling continous actions is relatively easy.\n",
        "- As its name suggests, in policy gradient we are following gradients with respect to the policy itself, which means we are constantly improving the policy. By contrast, in Q-learning we are improving our estimates of the values of different actions, which only implicitly improves the policy. You would think that improving the policy directly would be more efficient, and indeed it very often is.\n",
        "\n",
        "The main topic of this tutorial is policy gradient, covered in sections 1 through 3. In section 4, we will briefly go over another paradigm for solving reinforcement learning problems, namely model-based RL. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oAGF8RXcoRMN"
      },
      "source": [
        "---\n",
        "# Section 1: Intuition of Policy Gradient"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o_7WZG5nhcQf"
      },
      "source": [
        "Before learning the canonical policy gradient algorithm, let's build our intuition of policy optimization. Let $\\pi_{\\theta}$ be  a policy parametrized by $\\theta$, where $\\pi_{\\theta}(a|s)$ is the probability of taking an action $a$ in the current state $s$. Our goal is to iteratively update the parameters $\\theta$ such that $\\pi_{\\theta}$ eventually becomes the optimal policy. In particular, we want to find an update rule from $\\theta_t$ to $\\theta_{t+1}$ that leads us to the optimal policy $\\theta^*$.\n",
        "\n",
        "For now, let's assume discrete actions, where the output of the policy network is a probability distribution over all actions. To generalize to continuous actions, we can have the network parametrize a Gaussian distribution, similar to the encoder of VAE. In either case, we sample an action from the distribution to execute in inference time."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yHiDsdh7GYUL",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Run me (helper function)\n",
        "def simulate_policy_gradient(update_fn, filename, init=[0.0]*3):\n",
        "\n",
        "    # a,b,c are widgets which contain four variables\n",
        "    a = widgets.FloatSlider(min=0.0, max=1.0, value=1.0, step=0.010)\n",
        "    b = widgets.FloatSlider(min=0.0, max=1.0, value=1.0, step=0.005)\n",
        "    c = widgets.FloatSlider(min=0.0, max=1.0, value=1.0, step=0.0025)\n",
        "    sliders = [a, b, c]\n",
        "    \n",
        "    # initialize the logit and q-value for each action\n",
        "    for i, s in zip(init, sliders):\n",
        "        s.logit = i\n",
        "        s.q_val = s.step * 1000\n",
        "    \n",
        "    def update_values():\n",
        "        exps = [np.exp(e.logit) for e in sliders]\n",
        "        for ex, slid in zip(exps, sliders):\n",
        "            slid.value = ex / np.sum(exps)\n",
        "            slid.grad = slid.value * (1-slid.value)\n",
        "            \n",
        "    update_values()\n",
        "    np.random.seed(1)\n",
        "    it = 0\n",
        "    animation_data = []\n",
        "    while all([v.value < 0.95 for v in [a, b, c]]):\n",
        "        incr_idx, incr_size = update_fn(sliders)\n",
        "        update_values()\n",
        "        it += 1\n",
        "        animation_data.append((incr_idx, incr_size, [v.value for v in sliders]))\n",
        "        \n",
        "    if filename is not None:\n",
        "        fig, ax = plt.subplots()\n",
        "\n",
        "        def plot_animation(i):\n",
        "            plt.clf()\n",
        "            incr_idx, incr_size, values = animation_data[i]\n",
        "            plt.bar(['Val = %s' % (v.step * 1000) for v in sliders], values, width=0.5)\n",
        "            plt.ylim(0, 1.19)\n",
        "            if incr_size > 0:\n",
        "                plt.annotate('',\n",
        "                    xy=(incr_idx, values[incr_idx] + incr_size + 0.05), xycoords='data',\n",
        "                    xytext=(incr_idx, values[incr_idx]), textcoords='data',\n",
        "                    arrowprops=dict(width=5, connectionstyle=\"arc3\", color='green'),\n",
        "                )\n",
        "            else:\n",
        "                plt.annotate('',\n",
        "                    xy=(incr_idx, values[incr_idx]), xycoords='data',\n",
        "                    xytext=(incr_idx, values[incr_idx] - incr_size + 0.05), textcoords='data',\n",
        "                    arrowprops=dict(width=5, connectionstyle=\"arc3\", color='red'),\n",
        "                )\n",
        "            return fig,\n",
        "\n",
        "        ani = animation.FuncAnimation(fig, plot_animation, frames=list(range(0, len(animation_data))), blit=False)\n",
        "        ani.save(filename, writer='imagemagick', fps=10)\n",
        "        plt.close()\n",
        "    \n",
        "    print('Done in %s iterations' % it)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mBvjyb53qWB7"
      },
      "source": [
        "Let's look at a toy example to better understand how the policy optimization is done. In this toy example, there is a single state $s$ and three actions $a_1, a_2, a_3$. These actions have pre-defined Q-values `[10, 5, 2.5]`. We don't use a neural network yet, but assign a logit to each action. The policy is derived by taking the softmax of these logits. Our goal is to update the logits (which we denote $\\theta$) using gradient method so that the policy becomes optimal. \n",
        "\n",
        "\n",
        "Let $a^*$ be the optimal action. We want $\\pi_{\\theta}(a^*|s)$ to approach 1. If we had known the optimal action $a^*$, we can update the logits of $a^*$ by **gradient ascent**:\n",
        "\n",
        "$$\n",
        "\\theta_{t+1}=\\theta_{t}+\\alpha \\nabla \\pi_{\\theta_{t}}\\left(a^{*} | s\\right)\n",
        "$$\n",
        "\n",
        "Run the following cells to visualize how the results would like if we had known the optimal action $a^* = a_1$. The height of each bar is the probability of taking each action, and the arrow shows the gradient that we are following in each iteration."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A6buJGCa4_Vg",
        "cellView": "form"
      },
      "source": [
        "#@title Known optimal action\n",
        "\n",
        "lr = 0.5\n",
        "def update_known_optimal(actions):\n",
        "    # Each action has a .logit attribute that we must update. It also has a .grad attribute\n",
        "    # that gives the gradient of the that action.\n",
        "    actions[0].logit += lr * actions[0].grad\n",
        "    # We return the index of the action we updated and the size of the arrow to draw\n",
        "    # (which should be the size of our update, divided by a scaling factor to make the \n",
        "    # visual more appealing)\n",
        "    return 0, lr / 5\n",
        "    \n",
        "simulate_policy_gradient(update_known_optimal, 'known_optimal.gif')\n",
        "gifPath = Path(\"known_optimal.gif\")\n",
        "with open(gifPath,'rb') as f:\n",
        "    Image(data=f.read(), format='png')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rp5UWTe9o3a-"
      },
      "source": [
        "In practice, we do not know which actions are good. One solution would be to push on actions **proportional to our guess of the values** of these actions. These guesses can be highly approximate, but as long as they are somewhat grounded in reality, more overall pushing will happen on the optimal action $a^∗$. This way, it is guaranteed that our policy will converge to $\\pi_{\\theta}(a^∗|s)=1$.\n",
        "\n",
        "In the following demo, we sample an action $a$ uniformly at random in each iteration of policy update. We cheat a little and construct our value estimates by adding noise to the true value. Let $\\hat{Q} (s,a)$ denote our value estimate of action $a$, we apply the following update to the logits:\n",
        "\n",
        "$$\n",
        "\\theta_{t+1}=\\theta_{t}+\\alpha \\hat{Q}(s, a) \\nabla \\pi_{\\theta_{t}}(a | s)\n",
        "$$\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wzGkPS3ZGeKz",
        "cellView": "form"
      },
      "source": [
        "#@title Random action selection\n",
        "\n",
        "lr = 0.1\n",
        "def update_q_value(actions):\n",
        "    i = np.random.choice(list(range(len(actions))))\n",
        "    value_est = actions[i].q_val + np.random.randn() * actions[i].q_val\n",
        "    actions[i].logit += lr * value_est * actions[i].grad\n",
        "    return i, lr * value_est / 10\n",
        "    \n",
        "simulate_policy_gradient(update_q_value, 'random_action.gif')\n",
        "gifPath = Path('random_action.gif')\n",
        "with open(gifPath,'rb') as f:\n",
        "    Image(data=f.read(), format='png')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5yzS9B_hlowM"
      },
      "source": [
        "\n",
        "Here we assumed that we are given the values for each action. But in reality we do not know the values before training. How is the value generally computed (recollect what you learned in W11)?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tP0W5CdYuva1",
        "cellView": "form"
      },
      "source": [
        "value_estimate = \"\" #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z1IwBLnQI9Pb"
      },
      "source": [
        "# to_remove solution\n",
        "# Value is essentially the expectation of discounted rewards. We can rollout the \n",
        "# policy many times to termination and take the average discounted rewards."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ueGX68-icGfR"
      },
      "source": [
        "In practice, our agent is not going to choose actions uniformly. Rather, the actions are selected from the same policy $\\pi_{\\theta}$ that we are trying to learn. This is called **on-policy** training. There are two reasons why we might want to train on-policy:\n",
        "\n",
        "- We accumulate more rewards even as we train, which is desirable in some contexts.\n",
        "- It allows us to explore more promising areas of the state space by acting according to our current guess of the optimal actions.\n",
        "\n",
        "Let’s illustrate this phenomenon by using the same update rule as above but sampling actions according to their probabilities given by the current policy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RxVVBZzzGgYT",
        "cellView": "form"
      },
      "source": [
        "#@title On-Policy gradient update\n",
        "\n",
        "lr = 0.1\n",
        "def update_q_value_bad(actions):\n",
        "    # p decides the probability with each action is selected.\n",
        "    # Initially they are equal, but as value corresponding to one action\n",
        "    # becomes larger and larger, that action is selected more often\n",
        "    i = np.random.choice(list(range(len(actions))),p=[e.value for e in actions])\n",
        "    q_value_est = actions[i].q_val + np.random.randn() * actions[i].q_val\n",
        "    actions[i].logit += lr * q_value_est * actions[i].grad\n",
        "    return i, lr * q_value_est / 5\n",
        "\n",
        "init_logits = [0., 0., 1.5] \n",
        "simulate_policy_gradient(update_q_value_bad, 'bad.gif', init_logits)\n",
        "gifPath = Path('bad.gif')\n",
        "with open(gifPath,'rb') as f:\n",
        "    Image(data=f.read(), format='png')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dM8Z1DCPmjcb"
      },
      "source": [
        "\n",
        "In the above example we see that action 3 is selected despite action 1 being the better action. Why do you think so? Take a look at the `init_logits` in the above cell. For reference, the update rule is $\\theta_{t+1}=\\theta_{t}+\\alpha \\hat{Q}(s, a) \\nabla \\pi_{\\theta_{t}}(a \\mid s)$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "0IjA8WdXK3Xk"
      },
      "source": [
        "bad_actions = \"\" #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yb_9cZqXK_RY"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "# Although we update policy on actions with better values, we also end up pushing \n",
        "# more often on whichever actions happen to have higher values of $π_θ$ to begin \n",
        "# with (which could happen due to chance or bad initialization). These actions \n",
        "# might end up winning the race to the top in spite of being bad."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w-3dJ8qOoMRs"
      },
      "source": [
        "The above example shows that we need to compensate for the fact that more probable actions are going to be sampled more often. This leads us onto **on-policy correction**. The idea is to divide our update by the probability of the action. For example, if an action is 4x more likely to be taken than another, we will have 4x more gradient updates but each update would become 4x smaller. This is gives rise to the following update:\n",
        "\n",
        "$$\n",
        "\\theta_{t+1}=\\theta_{t}+\\alpha \\frac{\\hat{Q}(s, a) \\nabla \\pi_{\\theta_{t}}(a \\mid s)}{\\pi_{\\theta_t}(a \\mid s)}\n",
        "$$\n",
        "\n",
        "In this exercise, you will follow the instructions to fill in the gradient update rule for on-policy correction. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wbCBial55iAV"
      },
      "source": [
        "### Exercise 1: On-policy Correction"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QW8GhRVoTIPK"
      },
      "source": [
        "lr = 0.1\n",
        "def update_q_value_fixed(actions):\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Update Q value fixed\")\n",
        "    ####################################################################\n",
        "\n",
        "    # p (policy) decides the probability with each action is selected. As the\n",
        "    # probability of one action becomes larger, that action is selected more often.\n",
        "    # The probability of an action is stored in actions[i].value\n",
        "    num_actions = len(actions)\n",
        "    probs = [e.value for e in actions]\n",
        "    i = np.random.choice(list(range(num_actions)), p=probs)\n",
        "\n",
        "    # Noise is added to simulate values we get in training which vary with every episode\n",
        "    q_value_est = actions[i].q_val + np.random.randn() * actions[i].q_val\n",
        "\n",
        "    # Complete the update rule. alpha is given by lr. The gradient and probabiliy \n",
        "    # of an action can be accessed via actions[i].grad and actions[i].value.\n",
        "    actions[i].logit += ...\n",
        "\n",
        "    return i, lr * q_value_est / actions[i].value / 10\n",
        "\n",
        "# Uncomment on completion\n",
        "# simulate_policy_gradient(update_q_value_fixed, 'final_policy_grad.gif', init=[0.0, 0.0, 1.5])\n",
        "\n",
        "# gifPath = Path('final_policy_grad.gif')\n",
        "# with open(gifPath,'rb') as f:\n",
        "#     Image(data=f.read(), format='png')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4L2ovJP8o-xT"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "lr = 0.1\n",
        "def update_q_value_fixed(actions):\n",
        "    # p (policy) decides the probability with each action is selected. As the\n",
        "    # probability of one action becomes larger, that action is selected more often.\n",
        "    # The probability of an action is stored in actions[i].value\n",
        "    num_actions = len(actions)\n",
        "    probs = [e.value for e in actions]\n",
        "    i = np.random.choice(list(range(num_actions)), p=probs)\n",
        "\n",
        "    # Noise is added to simulate values we get in training which vary with every episode\n",
        "    q_value_est = actions[i].q_val + np.random.randn() * actions[i].q_val\n",
        "\n",
        "    # Complete the update rule. alpha is given by lr. The gradient and probabiliy \n",
        "    # of an action can be accessed via actions[i].grad and actions[i].value.\n",
        "    actions[i].logit += lr * q_value_est * actions[i].grad / actions[i].value\n",
        "    \n",
        "    return i, lr * q_value_est / actions[i].value / 10\n",
        "    \n",
        "simulate_policy_gradient(update_q_value_fixed, 'final_policy_grad.gif', init=[0.0, 0.0, 1.5])\n",
        "\n",
        "gifPath = Path('final_policy_grad.gif')\n",
        "with open(gifPath,'rb') as f:\n",
        "    Image(data=f.read(), format='png')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-eFiY164t6ca"
      },
      "source": [
        "---\n",
        "# Section 2: REINFORCE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UkcqCYx898rs",
        "cellView": "form"
      },
      "source": [
        "#@title Video: The REINFORCE Algorithm\n",
        "\n",
        "try: t1;\n",
        "except NameError: t1=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"jBRFiMWCCRI\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KK-Iwq9aHCuy"
      },
      "source": [
        "## Section 2.1: The REINFORCE algorithm\n",
        "\n",
        "Now that we've built some intuition on policy gradient using the toy example, we proceed to derive policy gradient for real RL problems. Recall that the reinforcement learning objective is to maximize the expect cumulative reward. Omitting discount factors for now, the objective can be written as follows: \n",
        "$$ J(\\theta) = E_{\\pi_{\\theta}}\\left[\\sum_{t=0}^{T}r(s_t, a_t)\\right]$$ \n",
        "\n",
        "Can we directly differentiate $J(\\theta)$ with respect to policy parameters $\\theta$? This isn't straightforward because the objective involves an expectation. In the VAE tutorial, we've seen one way of estimating gradients through an expectation: the reparametrization trick. Here we will use another technique called score function estimator. The full derivation of policy gradient using the score function estimator can be found in the appendix, but for our purposes we only need to know that $\\nabla_{\\theta}J(\\theta)$ takes the following form:\n",
        "\n",
        "\n",
        "\\begin{align*}\n",
        "\\nabla_{\\theta}J(\\theta) &= E_{\\tau \\sim p_{\\theta}(\\tau)} \\left[\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)\\sum_{t=1}^{T}r(s_t, a_t)\\right] \\\\\n",
        "&\\approx \\frac{1}{N} \\sum_{i=1}^{N} \\left(\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)\\sum_{t=1}^{T}r(s_t, a_t)\\right)\n",
        "\\end{align*}\n",
        "\n",
        "In the first expression above, $\\tau = (s_1, a_1, \\dots s_T, a_T)$ denotes a trajectory, and $p_{\\theta}(\\tau)$ is the distribution over trajectories induced by the policy. To estimate the policy gradient, we sample $N$ policy rollouts of $T$ steps each and compute the policy gradient using the second expression. This procedure is called **Monte-Carlo sampling**. Intuitively, we push up the log-probability of each action in proportion to the total obtained reward. In other words, we encourage actions seen in good trajectories and penalize those seen in bad trajectories.\n",
        "\n",
        "The vanilla policy gradient can suffer from a high variance. The issue is intrinsic to the probabilistic nature of the objective and Monte-Carlo sampling. As we unroll the policy forward in time, any small perturbation can lead to vastly different trajectories. So the average gradient could be uninformative. We could even have confilcting gradients because the same state-action pair may appear in both good and bad trajectories. The variance of gradients is also sensitive to reward magnitude. To see this, interact with the following demo and adjust `reward_weight` to see its impact. `reward_weight` is a constant used to scale the rewards.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YR262Hm4NFOF",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Demo: Variance\n",
        "\n",
        "# Report magnitudes of variances with different reward weights.\n",
        "\n",
        "from torch.distributions.multivariate_normal import MultivariateNormal \n",
        "\n",
        "reward_weight =  1#@param {type:\"number\"}\n",
        "act_dim = 2\n",
        "\n",
        "\n",
        "class Env:\n",
        "    def __init__(self, rew_wt):\n",
        "        self.goal = torch.tensor([10.0, 10.0])\n",
        "        self.rew_wt = rew_wt\n",
        "\n",
        "    def reward(self, state):\n",
        "        return - self.rew_wt * torch.sum((self.goal - state)**2)\n",
        "\n",
        "    def step(self, state, action):\n",
        "        next_state = state + action\n",
        "        reward = self.reward(next_state)\n",
        "        return next_state, reward\n",
        "\n",
        "\n",
        "def policy(state, theta):\n",
        "    mean = torch.sigmoid(state.T @ theta)\n",
        "    dist = MultivariateNormal(mean, covariance_matrix=0.001*torch.eye(act_dim))\n",
        "    action = dist.sample((1,))\n",
        "    return dist, action.squeeze()\n",
        "\n",
        "\n",
        "def demo_variance(rew_wt, N=50, T=50):\n",
        "    # Sample N rollouts of T steps each\n",
        "    env = Env(rew_wt)\n",
        "    init_state = torch.zeros((2))\n",
        "    theta = torch.zeros((act_dim, act_dim), requires_grad=True)\n",
        "    grads = []\n",
        "    for i in range(N):\n",
        "        state = init_state.clone()\n",
        "        log_likes, rewards = 0, 0\n",
        "        for t in range (T):\n",
        "            dist, action = policy(state, theta)\n",
        "            state, reward = env.step(state, action)\n",
        "            log_likes += dist.log_prob(action)\n",
        "            rewards += reward\n",
        "        J = log_likes * rewards\n",
        "        J.backward()\n",
        "        grad = theta.grad.clone()\n",
        "        grads.append(grad)\n",
        "    var_grads = torch.var(torch.stack(grads, 0), 0)\n",
        "    print(f\"Variance of gradients: \\n{var_grads.numpy()}\")\n",
        "    \n",
        "demo_variance(rew_wt=reward_weight)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "F1SgsdT90Iml"
      },
      "source": [
        "#@markdown Report the variance magnitude (e.g. 10^8) at these reward scales: 1, 1e-3, 1e-6.\n",
        "\n",
        "grad_vars = \"\" #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Wl9bidE7ob9W"
      },
      "source": [
        "We will see more ways to address the high variance issue later in this tutorial. For now, we can use two simple tricks to reduce the variance:\n",
        "\n",
        "1. Note that the action at time $t$ only influences future rewards. Instead of weighting gradients of all actions using the sum of the entire trajectory's rewards, we should weight them by the sum of future rewards (reward-to-go). \n",
        "2. We can discount future rewards to reduce the reward magnitude. This means we multiply the reward at time $t' \\geq t$ by $\\gamma^{t'-t}$.\n",
        "\n",
        "Applying these tricks, the policy gradient becomes \n",
        "\\begin{align*}\n",
        "\\nabla_{\\theta}J(\\theta) \\approx E_{\\tau \\sim p_{\\theta}(\\tau)} \\left[\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)\\sum_{t'=t}^{T}\\gamma^{t'-t}r(s_{t'}, a_{t'})\\right]\n",
        "\\end{align*}\n",
        "\n",
        "\n",
        "\n",
        "where the discounted future reward term $\\sum_{t'=t}^{T}\\gamma^{t'-t}r(s_{t'}, a_{t'})$ is often referred to as $G_t$. This policy gradient gives rise to the canonical algorithm called **REINFORCE**. Yes, you read it right, the name of the algorithm is exactly the name of the problem it's trying to solve. The pseudocode for REINFORCE is provided as follows\n",
        "\n",
        "Repeat:\n",
        "> 1. Sample trajectories $\\{\\tau_i\\}_{i=1}^N$ from the environment using the current policy. If $N=1$, then we update the policy at the end of each episode.\n",
        "2. Compute policy gradient: $\\nabla_{\\theta}J(\\theta) \\approx \\frac{1}{N} \\sum_{i=1}^{N} \\left(\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)G_t\\right)$\n",
        "3. Update policy: $\\theta_{t+1} \\leftarrow \\theta_{t} + \\alpha \\nabla_{\\theta} J(\\theta)$\n",
        "\n",
        "REINFORCE is an on-policy method. That is, in each iteration we must evaluate the gradient on new trajectories collected with the current policy. On-policy methods are expensive to train since they require large amounts of environment interactions. There are off-policy policy gradient methods which we will not introduce in this tutorial."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W08UO5rq5aNX"
      },
      "source": [
        "### Exercise 2 : Solving Cartpole"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yVU2lnwuocAY"
      },
      "source": [
        "To understand the REINFORCE algorithm better, lets use it to solve the [CartPole](https://gym.openai.com/envs/CartPole-v0/) environment from OpenAI Gym. In CartPole, a pole is attached by an un-actuated joint to a cart, which moves along a frictionless track. The system is controlled by applying a force of +1 or -1 to the cart. The pole starts upright, and the goal is to prevent it from falling over. A reward of +1 is provided for every timestep that the pole remains upright. An episode ends when the pole is more than 15 degrees from vertical, or the cart moves more than 2.4 units from the center. \n",
        "\n",
        "\n",
        "In this exercise, you will implement the REINFORCE algorithm. In particular, you will need to complete the `get_action` method in the policy network, which samples an action from the current policy, as well as `update_policy`, which performs one gradient update of the policy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3GO2hyjAsejR"
      },
      "source": [
        "class PolicyNetwork(nn.Module):\n",
        "    def __init__(self, num_inputs, num_actions, hidden_size, learning_rate=3e-4):\n",
        "        super(PolicyNetwork, self).__init__()\n",
        "        self.num_actions = num_actions\n",
        "        self.linear1 = nn.Linear(num_inputs, hidden_size)\n",
        "        self.linear2 = nn.Linear(hidden_size, num_actions)\n",
        "        self.optimizer = optim.Adam(self.parameters(), lr=learning_rate)\n",
        "\n",
        "    def forward(self, state):\n",
        "        x = F.relu(self.linear1(state))\n",
        "        x = F.softmax(self.linear2(x), dim=0)\n",
        "        return x \n",
        "\n",
        "    def get_action(self, state):\n",
        "        state = torch.from_numpy(state).to(device).float()\n",
        "        probs = self.forward(state)\n",
        "\n",
        "        ####################################################################\n",
        "        # Fill in missing code below (...),\n",
        "        # then remove or comment the line below to test your function\n",
        "        raise NotImplementedError(\"REINFORCE Policy\")\n",
        "        ####################################################################\n",
        "\n",
        "        # Sample an action from the policy. Hint: use np.random.choice()\n",
        "        action = ... # TODO\n",
        "\n",
        "        # Compute log probabilty for this action\n",
        "        log_prob = ... # TODO\n",
        "        return action, log_prob\n",
        "\n",
        "\n",
        "def update_policy(policy_network, rewards, log_probs, gamma):\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"REINFORCE Update\")\n",
        "    ####################################################################\n",
        "\n",
        "    discounted_rewards = []\n",
        "    for t in range(len(rewards)):\n",
        "        # At each step, we compute the sum of discounted future rewards\n",
        "        Gt = 0  # Gt is the sum of discounted future rewards\n",
        "        pow = 0 # pow keeps track of the power of discount factor gamma\n",
        "        for r in rewards[t:]:\n",
        "            # Compute discounted reward at the current time step\n",
        "            # Note: rewards is a list consisting of reward in each step of episode\n",
        "            Gt = Gt + ... # TODO\n",
        "            pow = pow + 1\n",
        "        discounted_rewards.append(Gt)\n",
        "    discounted_rewards = torch.tensor(discounted_rewards)\n",
        "\n",
        "    policy_gradients = []\n",
        "    for log_prob, Gt in zip(log_probs, discounted_rewards):\n",
        "        # Multiply log_prob with Gt and append to list of gradients.\n",
        "        # How can we have the optimizer perform gradient ASCENT?\n",
        "        policy_gradients.append(...) # TODO\n",
        "    \n",
        "    # Perform one step of gradient update\n",
        "    policy_network.optimizer.zero_grad()\n",
        "    objective = torch.stack(policy_gradients).sum()\n",
        "    objective.backward()\n",
        "    policy_network.optimizer.step()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X3iZ0dpGsrZd"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "class PolicyNetwork(nn.Module):\n",
        "    def __init__(self, num_inputs, num_actions, hidden_size, learning_rate=3e-4):\n",
        "        super(PolicyNetwork, self).__init__()\n",
        "        self.num_actions = num_actions\n",
        "        self.linear1 = nn.Linear(num_inputs, hidden_size)\n",
        "        self.linear2 = nn.Linear(hidden_size, num_actions)\n",
        "        self.optimizer = optim.Adam(self.parameters(), lr=learning_rate)\n",
        "\n",
        "    def forward(self, state):\n",
        "        x = F.relu(self.linear1(state))\n",
        "        x = F.softmax(self.linear2(x), dim=0)\n",
        "        return x \n",
        "\n",
        "    def get_action(self, state):\n",
        "        state = torch.from_numpy(state).to(device).float()\n",
        "        probs = self.forward(state)\n",
        "\n",
        "        # Sample an action from the policy. Hint: use np.random.choice()\n",
        "        action = np.random.choice(self.num_actions, p=probs.cpu().detach().numpy())\n",
        "\n",
        "        # Compute log probabilty for this action\n",
        "        log_prob = torch.log(probs[action])\n",
        "        return action, log_prob\n",
        "\n",
        "\n",
        "def update_policy(policy_network, rewards, log_probs, gamma):\n",
        "    discounted_rewards = []\n",
        "    for t in range(len(rewards)):\n",
        "        # At each step, we compute the sum of discounted future rewards\n",
        "        Gt = 0  # Gt is the sum of discounted future rewards\n",
        "        pow = 0 # pow keeps track of the power of discount factor gamma\n",
        "        for r in rewards[t:]:\n",
        "            # Compute discounted reward at the current time step\n",
        "            # Note: rewards is a list consisting of reward in each step of episode\n",
        "            Gt = Gt + gamma**pow * r\n",
        "            pow = pow + 1\n",
        "        discounted_rewards.append(Gt)\n",
        "    discounted_rewards = torch.tensor(discounted_rewards)\n",
        "\n",
        "    policy_gradients = []\n",
        "    for log_prob, Gt in zip(log_probs, discounted_rewards):\n",
        "        # Multiply log_prob with Gt and append to list of gradients.\n",
        "        # How can we have the optimizer perform gradient ASCENT?\n",
        "        policy_gradients.append(-log_prob * Gt)\n",
        "    \n",
        "    # Perform one step of gradient update\n",
        "    policy_network.optimizer.zero_grad()\n",
        "    objective = torch.stack(policy_gradients).sum()\n",
        "    objective.backward()\n",
        "    policy_network.optimizer.step()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61ykOFgPpfGD"
      },
      "source": [
        "Run the following cell to test your code. Training should take about 3 minutes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ub6Wn5r9pTEn"
      },
      "source": [
        "# Reinforce algorithm\n",
        "def reinforce(env, policy):\n",
        "    max_episode_num = 3000\n",
        "    max_steps = 10000\n",
        "    numsteps = []\n",
        "    avg_numsteps = []\n",
        "    all_rewards = []\n",
        "\n",
        "    for episode in tqdm(range(max_episode_num), position=0, leave=True):\n",
        "        state = env.reset()\n",
        "        log_probs = []\n",
        "        rewards = []\n",
        "\n",
        "        for steps in range(max_steps):\n",
        "            # Sample an action from the policy\n",
        "            action, log_prob = policy_net.get_action(state)\n",
        "            # Execute action in environment\n",
        "            new_state, reward, done, _ = env.step(action)\n",
        "            # Store log probability of action\n",
        "            log_probs.append(log_prob)\n",
        "            # Store reward\n",
        "            rewards.append(reward)\n",
        "\n",
        "            if done:\n",
        "                # Train policy at the end of an episode\n",
        "                update_policy(policy_net, rewards, log_probs, gamma=0.9)\n",
        "                numsteps.append(steps)\n",
        "                avg_numsteps.append(np.mean(numsteps[-10:]))\n",
        "                all_rewards.append(np.sum(rewards))\n",
        "                break\n",
        "            state = new_state\n",
        "        \n",
        "    plt.plot(numsteps)\n",
        "    plt.plot(avg_numsteps)\n",
        "    plt.xlabel('Episode')\n",
        "    plt.ylabel(\"Reward\")\n",
        "    plt.show()\n",
        "\n",
        "env = gym.make('CartPole-v0')\n",
        "policy_net = PolicyNetwork(env.observation_space.shape[0], env.action_space.n, 128)\n",
        "policy_net = policy_net.to(device)\n",
        "reinforce(env, policy_net)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gn7yBY69pXLE",
        "cellView": "form"
      },
      "source": [
        "# @markdown ## Run cell to generate video\n",
        "visualize(env, policy_net, 200)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B4PS0X8sQO5N"
      },
      "source": [
        "## Section 2.2: REINFORCE with whitening\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5TruvcxD4vCP"
      },
      "source": [
        "From the above result, we see that it takes a lot of episodes for the model to stabilize at a reward of 200. Sometimes we may also see a sharp decrease of reward during training. To further tackle the high variance problem of the REINFORCE algorithm, we can subtract a baseline $b(s_t)$ from $G_t$ while calculating the policy gradient: \n",
        "\n",
        "\\begin{align*}\n",
        "\\nabla_{\\theta}J(\\theta) &= E_{\\tau \\sim p_{\\theta}(\\tau)} \\left[\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)(G_t - b(s_t))\\right]\n",
        "\\end{align*}\n",
        "\n",
        "It can be shown that if the baseline only depends on $s_t$, then substracting the baseline still leads to an unbiased estimate of the policy gradient. See [this post](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html#baselines-in-policy-gradients) for more details. We will study a proper baseline in the next section, but there is a simple trick similar to a baseline which we can apply to greatly improve performance: we normalize the return at each step by the mean and standard deviation of all returns in the episode. This is called **whitening**. Whitening doesn't give us an unbiased estimator, but it works well in practice.\n",
        "\n",
        "\\begin{align*}\n",
        "G_{t}^{*} &= \\frac{G_{t}-\\bar{G}}{\\sigma_{G}} \\\\\n",
        "\\nabla_{\\theta}J(\\theta) &= E_{\\tau \\sim p_{\\theta}(\\tau)} \\left[\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)G_t^*\\right]\n",
        "\\end{align*}\n",
        "\n",
        "Rerun the below cell which implements a version of the REINFORCE algorithm with whitening."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A-bldTpy2Hza",
        "cellView": "form"
      },
      "source": [
        "# @markdown ##REINFORCE with whitening\n",
        "\n",
        "class PolicyNetwork(nn.Module):\n",
        "    def __init__(self, num_inputs, num_actions, hidden_size, learning_rate=3e-4):\n",
        "        super(PolicyNetwork, self).__init__()\n",
        "        self.num_actions = num_actions\n",
        "        self.linear1 = nn.Linear(num_inputs, hidden_size)\n",
        "        self.linear2 = nn.Linear(hidden_size, num_actions)\n",
        "        self.optimizer = optim.Adam(self.parameters(), lr=learning_rate)\n",
        "\n",
        "    def forward(self, state):\n",
        "        x = F.relu(self.linear1(state))\n",
        "        x = F.softmax(self.linear2(x), dim=1)\n",
        "        return x \n",
        "    \n",
        "    def get_action(self, state):\n",
        "        state = torch.from_numpy(state).float().unsqueeze(0)\n",
        "        probs = self.forward(state)\n",
        "        action = np.random.choice(self.num_actions, p=np.squeeze(probs.detach().numpy()))\n",
        "        log_prob = torch.log(probs.squeeze(0)[action])\n",
        "        return action, log_prob\n",
        "\n",
        "## Update\n",
        "def update_policy(policy_network, rewards, log_probs, gamma):\n",
        "    discounted_rewards = []\n",
        "    for t in range(len(rewards)):\n",
        "        Gt = 0 \n",
        "        pow = 0\n",
        "        for r in rewards[t:]:\n",
        "            Gt = Gt + gamma**pow * r\n",
        "            pow = pow + 1\n",
        "        discounted_rewards.append(Gt)\n",
        "    discounted_rewards = torch.tensor(discounted_rewards)\n",
        "    # Add whitening\n",
        "    discounted_rewards = (discounted_rewards - discounted_rewards.mean()) / (discounted_rewards.std() + 1e-9) # normalize discounted rewards\n",
        "\n",
        "    policy_gradients = []\n",
        "    for log_prob, Gt in zip(log_probs, discounted_rewards):\n",
        "        policy_gradients.append(-log_prob * Gt)\n",
        "    \n",
        "    policy_network.optimizer.zero_grad()\n",
        "    objective = torch.stack(policy_gradients).sum()\n",
        "    objective.backward()\n",
        "    policy_network.optimizer.step()\n",
        "\n",
        "## Main loop\n",
        "def reinforce_whitening(env, policy):\n",
        "    \n",
        "    max_episode_num = 3000\n",
        "    max_steps = 10000\n",
        "    numsteps = []\n",
        "    avg_numsteps = []\n",
        "    all_rewards = []\n",
        "\n",
        "    for episode in tqdm(range(max_episode_num),position=0, leave=True):\n",
        "        state = env.reset()\n",
        "        log_probs = []\n",
        "        rewards = []\n",
        "\n",
        "        for steps in range(max_steps):\n",
        "            action, log_prob = policy_net.get_action(state)\n",
        "            new_state, reward, done, _ = env.step(action)\n",
        "            log_probs.append(log_prob)\n",
        "            rewards.append(reward)\n",
        "\n",
        "            if done:\n",
        "                update_policy(policy_net, rewards, log_probs, 0.9)\n",
        "                numsteps.append(steps)\n",
        "                avg_numsteps.append(np.mean(numsteps[-10:]))\n",
        "                all_rewards.append(np.sum(rewards))\n",
        "                break\n",
        "            state = new_state\n",
        "        \n",
        "    plt.plot(numsteps)\n",
        "    plt.plot(avg_numsteps)\n",
        "    plt.ylabel(\"Reward\")\n",
        "    plt.xlabel('Episode')\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "env = gym.make('CartPole-v0')\n",
        "policy_net = PolicyNetwork(env.observation_space.shape[0], env.action_space.n, 128)\n",
        "policy_net = policy_net.to(device)\n",
        "reinforce_whitening(env, policy_net)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "cUoNR_GVJvG4"
      },
      "source": [
        "#@markdown ## Run cell to generate video\n",
        "visualize(env, policy_net, 200)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "_mDoNYqlU2UN"
      },
      "source": [
        "#@markdown What is the intuition behind whitening of expected rewards for reducing the variance between episodes?\n",
        "whitening_intuition = \"\" #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dIWldZ9eVFqI"
      },
      "source": [
        "# to_remove solution\n",
        "# Whitening reduces the variance of the gradient update by essentially decreasing\n",
        "# the magnitude of the discounted reward term. Smaller updates correspond to more \n",
        "# stable and smoother training."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G0igNmxSQQ4B"
      },
      "source": [
        "---\n",
        "# Section 3: Actor-Critic Algorithm"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TsZjYa7s-Iu1",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Actor-Critic Training\n",
        "\n",
        "try: t2;\n",
        "except NameError: t2=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"O8OEs5CfgYs\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NOH7DYgmhmDt"
      },
      "source": [
        "At the end of the last section, we breifly mentioned reducing the variance of policy gradient by subtracting a baseline from the discounted future reward at each step. What is a good candidate for the baseline? Is there a quantity that measures how good a state is **on average**? You may have guessed it: the value function! Indeed, a common choice of baseline is the value function $V(s)$, in which case the policy gradient becomes \n",
        "\n",
        "\\begin{align*}\n",
        "\\nabla_{\\theta}J(\\theta) &= E_{\\tau \\sim p_{\\theta}(\\tau)} \\left[\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)(G_t - V(s_t))\\right]\n",
        "\\end{align*}\n",
        "\n",
        "Here, the policy can be thought of as an actor, and the value function can be considered a critic that evaluates how good each action is. This is why we call it an **actor-critic algorithm**. During training, we alternately update the actor with policy gradient and the critic by minimizing the MSE between its prediction and the sum of discounted future rewards.\n",
        "\n",
        "Let's take a closer look at the expression. Recall that the value function essentially tells us how good a state is, averaged across all actions. So intuitively we encourage actions that do better than average, and penalize those that do worse. This algorithm is also called the **advantage actor critic (A2C)** because $G_t - V(s_t)$ approximates the advantage function of a state-ation pair defined as $$A(s_t, a_t) = Q(s_t, a_t) - V(s_t)$$\n",
        "\n",
        "From an advantage function point of view, we are approximating the Q-value of $s$ and $a$ using the discounted future reward of the current trajectory. To stablize training, we can further bootstrap an estimate of the Q-value from the value function (critic). Instead of executing all the way to termination, we rollout $n$-steps and compute $$Q(s_t, a_t) = \\sum_{t'=t}^{t+n-1}\\gamma^{t'-t}r(s_{t'}, a_{t'}) + \\gamma^{n}V(s_{t+n})$$\n",
        "\n",
        "That is, we take $n-1$ real rewards and bootstrap an expected sum of future rewards from the value function. This relates back to TD learning introduced in week 11. In particular, if $n=1$, then the advantage function is approximated by the TD error: $A(s,a) = r + \\gamma V(s') - V(s)$. We present an analysis of the actor-critic algorithm from a TD perspective in the appendix. \n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4f7KhQJs5dz4"
      },
      "source": [
        "## Exercise 3: Advantage actor-critic\n",
        "\n",
        "In the following exercise, you will implement the advantage actor-critic algorithm. Complete the main `a2c` function and train an actor-critic agent to solve the CartPole environment. Training should take about 7 minutes."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ahCoW3WxDQBs"
      },
      "source": [
        "# Actor network\n",
        "class Actor(nn.Module):\n",
        "    def __init__(self, state_dim, num_actions, learning_rate=1e-3):\n",
        "        super().__init__()\n",
        "        self.num_actions = num_actions\n",
        "        self.model = nn.Sequential(\n",
        "            nn.Linear(state_dim, 64),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(64, 32),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(32, num_actions),\n",
        "            nn.Softmax(0)\n",
        "        )\n",
        "        self.optimizer = optim.Adam(self.parameters(), lr=learning_rate)\n",
        "    \n",
        "    def forward(self, state):\n",
        "        return self.model(state)\n",
        "\n",
        "    def get_action(self, state):\n",
        "        state = torch.from_numpy(state).to(device).float()\n",
        "        probs = self.forward(state)\n",
        "        # Sample an action from the policy.\n",
        "        action = np.random.choice(self.num_actions, p=probs.cpu().detach().numpy())\n",
        "        # Compute log probabilty for this action\n",
        "        log_prob = torch.log(probs[action])\n",
        "        return action, log_prob\n",
        "\n",
        "# Critic network\n",
        "class Critic(nn.Module):\n",
        "    def __init__(self, state_dim, learning_rate=1e-3):\n",
        "        super().__init__()\n",
        "        self.model = nn.Sequential(\n",
        "            nn.Linear(state_dim, 64),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(64, 32),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(32, 1)\n",
        "        )\n",
        "        self.optimizer = optim.Adam(self.parameters(), lr=learning_rate)\n",
        "    \n",
        "    def forward(self, state):\n",
        "        return self.model(state)\n",
        "\n",
        "    def get_value(self, state):\n",
        "        return self.forward(torch.from_numpy(state).to(device).float())[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i9YJ7I4_aldU"
      },
      "source": [
        "def a2c(env, actor, critic, max_episodes, num_steps, gamma=0.99):\n",
        "\n",
        "    ####################################################################\n",
        "    # Fill in missing code below (...),\n",
        "    # then remove or comment the line below to test your function\n",
        "    raise NotImplementedError(\"Actor critic\")\n",
        "    ####################################################################\n",
        "\n",
        "    all_lengths = []\n",
        "    average_lengths = []\n",
        "    all_rewards = []\n",
        "\n",
        "    for episode in tqdm(range(max_episodes), position=0, leave=True):\n",
        "        log_probs = []\n",
        "        values = []\n",
        "        rewards = []\n",
        "\n",
        "        state = env.reset()\n",
        "        for steps in range(num_steps):\n",
        "            # Get the value of the current state from the critic\n",
        "            value = ... # TODO\n",
        "\n",
        "            # Get an action and its log probability from the actor\n",
        "            action, log_prob = ... # TODO\n",
        "\n",
        "            # Execute action in environment \n",
        "            new_state, reward, done, _ = env.step(action)\n",
        "\n",
        "            # Store value, log probability, and reward\n",
        "            values.append(value)\n",
        "            log_probs.append(log_prob)\n",
        "            rewards.append(reward)\n",
        "\n",
        "            # Update state\n",
        "            state = new_state\n",
        "            \n",
        "            if done or steps == num_steps - 1:\n",
        "                # Store value of the last state\n",
        "                V_last = critic.get_value(new_state)\n",
        "                all_rewards.append(np.sum(rewards))\n",
        "                break\n",
        "        \n",
        "        \n",
        "        values = torch.stack(values)\n",
        "        log_probs = torch.stack(log_probs)\n",
        "\n",
        "        # Compute Q values. Qvals[t] = r_t + gamma * r_{t+1} + ... + gamma^{T-t} * v_T\n",
        "        Qvals = torch.zeros_like(values)\n",
        "        # At this point, V_last stores the value of the last state, which we bootstrap from.\n",
        "        Qval = V_last\n",
        "        for t in reversed(range(len(rewards))):\n",
        "            Qval = ... # TODO\n",
        "            Qvals[t] = Qval\n",
        "        \n",
        "        \n",
        "        # Compute advantage for all steps. Use Q_vals and values.\n",
        "        advantage = ... # TODO\n",
        "\n",
        "        # Update critic\n",
        "        critic_loss = advantage.pow(2).mean()\n",
        "        critic.optimizer.zero_grad()\n",
        "        critic_loss.backward()\n",
        "        critic.optimizer.step()\n",
        "        \n",
        "        # Update critic\n",
        "        actor_loss = torch.mean(-log_probs * advantage.detach())\n",
        "        actor.optimizer.zero_grad()\n",
        "        actor_loss.backward()\n",
        "        actor.optimizer.step()\n",
        "\n",
        "    # Plot results\n",
        "    smoothed_rewards = pd.Series.rolling(pd.Series(all_rewards), 10).mean()\n",
        "    smoothed_rewards = [elem for elem in smoothed_rewards]\n",
        "    plt.plot(all_rewards)\n",
        "    plt.plot(smoothed_rewards)\n",
        "    plt.plot()\n",
        "    plt.xlabel('Episode')\n",
        "    plt.ylabel('Reward')\n",
        "    plt.show()\n",
        "\n",
        "# Uncomment on completion\n",
        "# # Environment\n",
        "# env = gym.make(\"CartPole-v0\")\n",
        "# state_dim = env.observation_space.shape[0]\n",
        "# n_actions = env.action_space.n\n",
        "# # Actor and critic models\n",
        "# actor = Actor(state_dim, n_actions).to(device)\n",
        "# critic = Critic(state_dim).to(device)\n",
        "# # Hyperparameters\n",
        "# max_episodes = 3000\n",
        "# num_steps = 300\n",
        "# a2c(env, actor, critic, max_episodes, num_steps)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F_bEsXsCZsmy"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "def a2c(env, actor, critic, max_episodes, num_steps, gamma=0.99):\n",
        "    all_lengths = []\n",
        "    average_lengths = []\n",
        "    all_rewards = []\n",
        "\n",
        "    for episode in tqdm(range(max_episodes), position=0, leave=True):\n",
        "        log_probs = []\n",
        "        values = []\n",
        "        rewards = []\n",
        "\n",
        "        state = env.reset()\n",
        "        for steps in range(num_steps):\n",
        "            # Get the value of the current state from the critic\n",
        "            value = critic.get_value(state)\n",
        "\n",
        "            # Get an action and its log probability from the actor\n",
        "            action, log_prob = actor.get_action(state)\n",
        "\n",
        "            # Execute action in environment \n",
        "            new_state, reward, done, _ = env.step(action)\n",
        "\n",
        "            # Store value, log probability, and reward\n",
        "            values.append(value)\n",
        "            log_probs.append(log_prob)\n",
        "            rewards.append(reward)\n",
        "\n",
        "            # Update state\n",
        "            state = new_state\n",
        "            \n",
        "            if done or steps == num_steps - 1:\n",
        "                # Store value of the last state\n",
        "                V_last = critic.get_value(new_state)\n",
        "                all_rewards.append(np.sum(rewards))\n",
        "                break\n",
        "        \n",
        "        \n",
        "        values = torch.stack(values)\n",
        "        log_probs = torch.stack(log_probs)\n",
        "\n",
        "        # Compute Q values. Qvals[t] = r_t + gamma * r_{t+1} + ... + gamma^{T-t} * v_T\n",
        "        Qvals = torch.zeros_like(values)\n",
        "        # At this point, V_last stores the value of the last state, which we bootstrap from.\n",
        "        Qval = V_last\n",
        "        for t in reversed(range(len(rewards))):\n",
        "            Qval = rewards[t] + gamma * Qval\n",
        "            Qvals[t] = Qval\n",
        "        \n",
        "        \n",
        "        # Compute advantage for all steps\n",
        "        advantage = Qvals - values\n",
        "\n",
        "        # Update critic\n",
        "        critic_loss = advantage.pow(2).mean()\n",
        "        critic.optimizer.zero_grad()\n",
        "        critic_loss.backward()\n",
        "        critic.optimizer.step()\n",
        "        \n",
        "        # Update critic\n",
        "        actor_loss = torch.mean(-log_probs * advantage.detach())\n",
        "        actor.optimizer.zero_grad()\n",
        "        actor_loss.backward()\n",
        "        actor.optimizer.step()\n",
        "\n",
        "    # Plot results\n",
        "    smoothed_rewards = pd.Series.rolling(pd.Series(all_rewards), 10).mean()\n",
        "    smoothed_rewards = [elem for elem in smoothed_rewards]\n",
        "    plt.plot(all_rewards)\n",
        "    plt.plot(smoothed_rewards)\n",
        "    plt.plot()\n",
        "    plt.xlabel('Episode')\n",
        "    plt.ylabel('Reward')\n",
        "    plt.show()\n",
        "\n",
        "# Environment\n",
        "env = gym.make(\"CartPole-v0\")\n",
        "state_dim = env.observation_space.shape[0]\n",
        "n_actions = env.action_space.n\n",
        "# Actor and critic models\n",
        "actor = Actor(state_dim, n_actions).to(device)\n",
        "critic = Critic(state_dim).to(device)\n",
        "# Hyperparameters\n",
        "max_episodes = 3000\n",
        "num_steps = 300\n",
        "a2c(env, actor, critic, max_episodes, num_steps)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "0vBHMCsTuYnA"
      },
      "source": [
        "# @markdown ## Run cell to generate video\n",
        "visualize(env, actor, 200)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j3qjGuR1-XCy",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Model-Free RL Examples\n",
        "\n",
        "video = YouTubeVideo(id=\"m_bU6gb9W4M\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JeG_hpR-oOsV"
      },
      "source": [
        "---\n",
        "# Section 4: Model-Based Reinforcement Learning"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P4MrNKhb-k3K",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Model-Based RL Intro\n",
        "\n",
        "try: t3;\n",
        "except NameError: t3=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"54q4Wf7anmc\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wFq4ayYo-q7k"
      },
      "source": [
        "One issue with RL algorthms we've seen so far is bad sampling efficiency. Both DQN and policy gradient require huge amounts of environment interactions to be effective. This makes it difficult to apply reinforcement learning to problems where data colloection can be slow or expensive. Sampling efficiency can be greatly improved if we have an internal representation of the environment from which we can obtain additional data. This technique is called model-based reinforcement learning (MBRL). In MBRL, we learn a model of the transition dynamics (and sometimes the reward function) of the environment. The simplist model is a neural network $f_{\\theta}$ which predicts the next state give the current state and action: $s_{t+1} = f_{\\theta}(s_t, a_t)$. Having a model opens up many possibilities. Broadly speaking, we can either augment policy DQN / gradient with experience sampled from the model, or apply planning techniques to find the optimal action.\n",
        "\n",
        "Besides good sampling efficiency, another motivation for MBRL is fast adaptation. Consider a task where a robotic arm needs to reach a goal. We can easily train the agent to reach a fixed goal, but what if we want it to reach another goal? Model-free RL techniques would have to learn everything from scratch again. However, note that when the goal changes, the underlying physics doesn't. Thus, using MBRL, we can keep the dynamics model and only adapt the reward model to the new goal, enabling the agent to quickly solve the new task."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3WvnlmkocyUs",
        "cellView": "form"
      },
      "source": [
        "#@markdown If we add an obstacle to the goal-reaching task, do we need to retrain the dynamics model? Why?\n",
        "obstacle = \"\" #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VY1AcLYYeRJd"
      },
      "source": [
        "# to_remove solution\n",
        "# Yes, we need to learn a new dynamics model because the robotic arm may hit the \n",
        "# obstacle when reaching the goal and our model wouldn't know how to handle it."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aQrk1ViG-5tK",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Model Learning and Usage\n",
        "\n",
        "video = YouTubeVideo(id=\"MmF8dYifZDY\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hh5CJx2aTyLP"
      },
      "source": [
        "## Section 4.1: Planning with world models\n",
        "\n",
        "We will not go into details of model-based policy optimization or Q-learning. Instead, we will dedicate this section to introducing some planning techniques that can be used to directly find the optimal action given a learned model. Given a model, here are a few things we can do:\n",
        "\n",
        "\n",
        "1. If the state space is small, we can apply dynamic programming techniques like policy/value iteration to find an optimal policy. Recall that policy/value iteration assumes full knowledge of the dynamics and rewards, which the model approximates.\n",
        "2. If the state space is large, we can consider an approach introduced all the way back in week 1: Monte-Carlo tree search. Note that AlphaZero is not model-based. We use MCTS because we know the transitions of an Chess game.\n",
        "3. The previous cases assume discrete actions. With a continuous action space, we resort to trajectory optimization techniques. \n",
        "\n",
        "\n",
        "Let's take a look at a sampling-based trajectory optimization technique called the cross-entropy (CE) method. Say we are currently at state $s$. Our goal is to find a sequence of actions $a_1, \\dots a_H$ that maximize the reward. $H$ is the horizon of the plan. To do this, let's try modeling the action squence using an isotropic gaussian distribution, with means $\\mu_1, \\dots, \\mu_H$ and standard deviations $\\sigma_1, \\dots \\sigma_H$. We want to update this distribution to eventually sample a good sequence from it. To do this, we initialize the distribution to unit gaussian and repeat the following process:\n",
        "```\n",
        "Repeat for s steps:\n",
        "    Sample n action sequences from the current distribution.\n",
        "    Roll out the sequences through our model (starting at s) and compute the reward for each trajectory.\n",
        "    Re-fit our distribution to the top k action sequences. We do this by updating the means and stds to the means and stds of these k samples.\n",
        "```\n",
        "Finally, we sample an action sequence from the final distribution and execute the plan. In practice, the model is not perfect, and its prediction error will compound as we roll out many steps. Thus we typically execute only the first action from the sampled sequence and replan. This is called model-predictive control (MPC).\n",
        "\n",
        "The following code is a demo for the CE method. We created a toy environment where a point is initialized at (0, 0) and needs to reach (10, 10). The actions are offsets in the x and y directions, and they are clipped to be within (-1, -1) and (1, 1). We use the environment itself as a perfect model and apply the CE method to plan from the initial state. \n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WisD2-TXdXvN"
      },
      "source": [
        "# Define action dimension\n",
        "act_dim = 2\n",
        "\n",
        "class PointEnv:\n",
        "    def __init__(self):\n",
        "        self.goal = torch.tensor([10.0, 10.0])\n",
        "        self.init_state = torch.tensor([0., 0.])\n",
        "\n",
        "    def reward(self, state):\n",
        "        return - 0.1 * torch.sum((self.goal - state)**2, dim=-1)\n",
        "\n",
        "    def step(self, state, action):\n",
        "        next_state = state + torch.clip(action, -1, 1)\n",
        "        reward = self.reward(next_state)\n",
        "        return next_state, reward\n",
        "\n",
        "    def unroll(self, state, actions):\n",
        "        num_samples, horizon = actions.shape[0:2]\n",
        "        states = state.unsqueeze(0).repeat(num_samples, 1)\n",
        "        rewards_total = torch.zeros(num_samples)\n",
        "        for t in range(horizon):\n",
        "            next_states, rewards = self.step(states, actions[:, t])\n",
        "            states = next_states\n",
        "            rewards_total += rewards\n",
        "        return rewards_total\n",
        "\n",
        "# CEM on point environment with known transitions and reward\n",
        "def cross_entropy_method(model, state, horizon, n=100, k=10, num_iters=10):\n",
        "    act_mean = torch.zeros((horizon, act_dim))\n",
        "    act_std = torch.ones((horizon, act_dim))\n",
        "    for i in range(num_iters):\n",
        "        # Create distribution using mean and std\n",
        "        act_dist = Normal(act_mean, act_std)\n",
        "        # Sample n action sequences\n",
        "        actions = act_dist.sample([n])\n",
        "        # Unroll squences through model and compute rewards\n",
        "        rewards = model.unroll(state, actions)\n",
        "        # Take the top k samples\n",
        "        top_rews, top_inds = torch.topk(rewards, k)\n",
        "        top_actions = actions[top_inds]\n",
        "        # Re-fit distribution\n",
        "        act_std, act_mean = torch.std_mean(top_actions, dim=0)\n",
        "    return act_mean, act_std\n",
        "\n",
        "# Use the environment as a perfect model\n",
        "model = PointEnv()\n",
        "init_state = model.init_state\n",
        "act_mean, act_std = cross_entropy_method(model, init_state, horizon=10)\n",
        "print(f\"Mean of actions:\\n{act_mean.numpy()}\\nStd of actions:\\n{act_std.numpy()}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ripwTUhJIqsg",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Run cell to execute an action sequence sampled from the planned distribution.\n",
        "\n",
        "# Sample an action sequence from action distribution and execute all of them\n",
        "# Our environment clips actions to within (-1, 1) so if all actions in our \n",
        "# sequence are greater than (1, 1) then we will reach the goal in 10 steps.\n",
        "def execute_cem(env, init_state, act_mean, act_std):\n",
        "    state = init_state\n",
        "    act_dist = Normal(act_mean, act_std)\n",
        "    actions = act_dist.sample()\n",
        "    for act in actions:\n",
        "        next_state, reward = env.step(state, act)\n",
        "        diff = next_state - state\n",
        "        plt.arrow(state[0], state[1], diff[0], diff[1], length_includes_head=True, head_width=0.15)\n",
        "        state = next_state\n",
        "    \n",
        "    dist = torch.sum((state - env.goal)**2)\n",
        "    print(f\"Final distance to goal: {dist}\")\n",
        "\n",
        "    plt.scatter(init_state[0], init_state[1])\n",
        "    plt.scatter(env.goal[0], env.goal[1])\n",
        "    plt.show()\n",
        "\n",
        "execute_cem(model, init_state, act_mean, act_std)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fpWx_xopaIhX"
      },
      "source": [
        "## Section 4.2: Approximating dynamics with a world model\n",
        "Model learning is in theory off-policy, meaning it can learn from any data. However, in practice MBRL will perform poorly if we don't train the model with on-policy data. This makes sense because the model has limited capacity and we want it be accurate on parts of the state space that we care about. Below is a sketch of an MBRL algorithm with on-policy data collection.\n",
        "\n",
        "```\n",
        "Run base policy (e.g. random policy) to collect a dataset D of transition tuples.\n",
        "Repeat\n",
        "    Train model using samples from D.\n",
        "    If there is a policy, update it with trajectories sampled from the model.\n",
        "    Collect new data from the environment by executing the policy or planning from the model. Add new data to D.\n",
        "```\n",
        "\n",
        "\n",
        "We will not implement a full MBRL algorithm in this tutorial. Instead, you will train a dynamics model for the CartPole environment on randomly collected data. This will look a lot like supervised learning. The only caveat is that you will train the model to predict the state difference $\\Delta s = s_{t+1} - s_{t}$, which is easier to learn when $s_t$ and $s_{t+1}$ are very similar. The next state can be reconstructed by adding the difference to the current state. Follow the instructions to fill in the `update` and `predict` methods of the dynamics model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FSd4L3YzfoHv"
      },
      "source": [
        "### Exercise 4: Train a dynamics model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3buaBqoqdni5"
      },
      "source": [
        "class DynamicsModel(nn.Module):\n",
        "    def __init__(self, act_dim, state_dim, hidden_size, learning_rate=0.001):\n",
        "        super().__init__()\n",
        "        self.act_dim = act_dim\n",
        "        self.state_dim = state_dim\n",
        "        self.learning_rate = learning_rate\n",
        "        \n",
        "        # 2-layer fully connected network\n",
        "        self.layers = nn.Sequential(\n",
        "            nn.Linear(act_dim + state_dim, hidden_size),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(hidden_size, hidden_size),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(hidden_size, state_dim)\n",
        "        )\n",
        "\n",
        "        # Use Adam optimizer and MSELoss()\n",
        "        self.optimizer = optim.Adam(self.layers.parameters(), self.learning_rate)\n",
        "        self.criterion = nn.MSELoss()\n",
        "    \n",
        "    def forward(self, input):\n",
        "        return self.layers(input)\n",
        "\n",
        "    def update(self, state, action, next_state):\n",
        "        # Update model\n",
        "        # state: [n, state_dim]\n",
        "        # action: [n, act_dim]\n",
        "        # next_state: [n, state_dim]\n",
        "\n",
        "        ####################################################################\n",
        "        # Fill in missing code below (...),\n",
        "        # then remove or comment the line below to test your function\n",
        "        raise NotImplementedError(\"Model update\")\n",
        "        ####################################################################\n",
        "\n",
        "        # Concatenate state and action along the last dimension\n",
        "        input = ...\n",
        "\n",
        "        # Forward input through model to get the output\n",
        "        output = ...\n",
        "\n",
        "        # Constuct target as the state difference \n",
        "        target = ...\n",
        "\n",
        "        # Compute model loss using self.criterion\n",
        "        loss = ...\n",
        "\n",
        "        # Take one gradient step using self.optimizer\n",
        "        self.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        self.optimizer.step()\n",
        "        return loss.item()\n",
        "    \n",
        "    def predict(self, state, action):\n",
        "        ####################################################################\n",
        "        # Fill in missing code below (...),\n",
        "        # then remove or comment the line below to test your function\n",
        "        raise NotImplementedError(\"Model predict\")\n",
        "        ####################################################################\n",
        "\n",
        "        # Concatenate state and action along the last dimension\n",
        "        input = ...\n",
        "\n",
        "        with torch.no_grad():\n",
        "            # Forward input through model to get the state difference\n",
        "            diff = ...\n",
        "            \n",
        "        # Add difference to input state to reconstruct the next state.\n",
        "        next_state = ...\n",
        "        return next_state"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B048DaXuhgTN"
      },
      "source": [
        "# to_remove solution\n",
        "\n",
        "class DynamicsModel(nn.Module):\n",
        "    def __init__(self, act_dim, state_dim, hidden_size, learning_rate=0.001):\n",
        "        super().__init__()\n",
        "        self.act_dim = act_dim\n",
        "        self.state_dim = state_dim\n",
        "        self.learning_rate = learning_rate\n",
        "        \n",
        "        # 2-layer fully connected network\n",
        "        self.layers = nn.Sequential(\n",
        "            nn.Linear(act_dim + state_dim, hidden_size),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(hidden_size, hidden_size),\n",
        "            nn.Tanh(),\n",
        "            nn.Linear(hidden_size, state_dim)\n",
        "        )\n",
        "\n",
        "        # Use Adam optimizer and MSELoss()\n",
        "        self.optimizer = optim.Adam(self.layers.parameters(), self.learning_rate)\n",
        "        self.criterion = nn.MSELoss()\n",
        "    \n",
        "    def forward(self, input):\n",
        "        return self.layers(input)\n",
        "\n",
        "    def update(self, state, action, next_state):\n",
        "        # Update model\n",
        "        # state: [n, state_dim]\n",
        "        # action: [n, act_dim]\n",
        "        # next_state: [n, state_dim]\n",
        "\n",
        "        # Concatenate state and action along the last dimension\n",
        "        input = torch.cat((state, action), -1)\n",
        "\n",
        "        # Forward input through model to get the output\n",
        "        output = self.forward(input)\n",
        "\n",
        "        # Constuct target as the state difference \n",
        "        target = next_state - state\n",
        "\n",
        "        # Compute model loss using self.criterion\n",
        "        loss = self.criterion(output, target)\n",
        "\n",
        "        # Take one gradient step using self.optimizer\n",
        "        self.optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        self.optimizer.step()\n",
        "\n",
        "        return loss.item()\n",
        "    \n",
        "    def predict(self, state, action):\n",
        "        # Concatenate state and action along the last dimension\n",
        "        input = torch.cat((state, action), -1)\n",
        "\n",
        "        with torch.no_grad():\n",
        "            # Forward input through model to get the state difference\n",
        "            diff = self.forward(input)\n",
        "\n",
        "        # Add difference to input state to reconstruct the next state.\n",
        "        next_state = state + diff\n",
        "        return next_state"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XrN8L2J1CGJt",
        "cellView": "form"
      },
      "source": [
        "#@markdown ## Run cell to train model\n",
        "\n",
        "class ReplayBuffer:\n",
        "  def __init__(self, state_dim, act_dim, buffer_size):\n",
        "    self.buffer_size = buffer_size\n",
        "    self.ptr = 0\n",
        "    self.n_samples = 0\n",
        "\n",
        "    self.state = torch.zeros(buffer_size, *state_dim, dtype=torch.float32, device=device)\n",
        "    self.action = torch.zeros(buffer_size, act_dim, dtype=torch.int64, device=device)\n",
        "    self.reward = torch.zeros(buffer_size, 1, dtype=torch.float32, device=device)\n",
        "    self.next_state = torch.zeros(buffer_size, *state_dim, dtype=torch.float32, device=device)\n",
        "\n",
        "  def add(self, state, action, reward, next_state):\n",
        "    self.state[self.ptr] = torch.tensor(state)\n",
        "    self.action[self.ptr] = torch.tensor(action)\n",
        "    self.reward[self.ptr] = torch.tensor(reward)\n",
        "    self.next_state[self.ptr] = torch.tensor(next_state)\n",
        "    \n",
        "    if self.n_samples < self.buffer_size:\n",
        "      self.n_samples += 1\n",
        "\n",
        "    self.ptr = (self.ptr + 1) % self.buffer_size\n",
        "\n",
        "  def sample(self, batch_size):      \n",
        "    # Select batch_size number of sample indicies at random from the buffer\n",
        "    idx = np.random.choice(self.n_samples, batch_size)    \n",
        "    # Using the random indices, assign the corresponding state, action, reward,\n",
        "    # and next state samples.\n",
        "    state = self.state[idx]\n",
        "    action = self.action[idx]\n",
        "    reward = self.reward[idx]\n",
        "    next_state = self.next_state[idx]\n",
        "    return state, action, reward, next_state\n",
        "\n",
        "def collect_random_data(env, num_samples=10000):\n",
        "    print('Collecting random data')\n",
        "    replay_buffer = ReplayBuffer(state_dim=(4,), act_dim=1, buffer_size=num_samples)\n",
        "    state = env.reset()\n",
        "    while replay_buffer.n_samples < replay_buffer.buffer_size:\n",
        "        # Sample a random action\n",
        "        action = env.action_space.sample()\n",
        "        # Execute action in environment\n",
        "        next_state, reward, done, info = env.step(action)\n",
        "        # Add to replay buffer\n",
        "        replay_buffer.add(state, action, reward, next_state)\n",
        "        # Update state\n",
        "        state = next_state\n",
        "        if done:\n",
        "            state = env.reset()\n",
        "    return replay_buffer\n",
        "\n",
        "\n",
        "def train_model(model, replay_buffer, num_iters, batch_size):\n",
        "    print('Training model on collected data')\n",
        "    losses = []\n",
        "    for i in tqdm(range(num_iters)):\n",
        "        # Sample batch from replay buffer\n",
        "        states, actions, rewards, next_states = replay_buffer.sample(batch_size)\n",
        "        # Update model\n",
        "        loss = model.update(states, actions, next_states)\n",
        "        losses.append(loss)\n",
        "\n",
        "    # Show loss plot\n",
        "    plt.plot(losses)\n",
        "    plt.title('Dynamics loss')\n",
        "    plt.xlabel('Iteration')\n",
        "    plt.ylabel('Loss')\n",
        "    plt.show()\n",
        "    \n",
        "\n",
        "def eval_model(model, env):\n",
        "    # Show prediction errors\n",
        "    state = env.reset()\n",
        "    for i in range(5):\n",
        "        action = env.action_space.sample()\n",
        "        next_state, _, _, _ = env.step(action)\n",
        "        state_t = torch.tensor(state, dtype=torch.float32, device=device)\n",
        "        action_t = torch.tensor([action], dtype=torch.float32, device=device)\n",
        "        next_state_pred = model.predict(state_t, action_t).cpu().numpy()\n",
        "        error = next_state - next_state_pred\n",
        "        print(f'Step {i}, prediction error: {error}')\n",
        "        state = next_state\n",
        "\n",
        "env = gym.make('CartPole-v0')\n",
        "replay_buffer = collect_random_data(env)\n",
        "model = DynamicsModel(act_dim=1, state_dim=4, hidden_size=10).to(device)\n",
        "train_model(model, replay_buffer, num_iters=1000, batch_size=50)\n",
        "eval_model(model, env)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PbnbGRdh_Fd2",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Model-Based RL Examples\n",
        "\n",
        "try: t4;\n",
        "except NameError: t4=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"ejouP7Qv3AA\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fjVyNSrsfDxB"
      },
      "source": [
        "---\n",
        "# Wrap-up"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R0rR2gmDfIM9",
        "cellView": "form"
      },
      "source": [
        "import time\n",
        "import numpy as np\n",
        "import urllib.parse\n",
        "from IPython.display import IFrame\n",
        "\n",
        "#@markdown #Run Cell to Show Airtable Form\n",
        "#@markdown ##**Confirm your answers and then click \"Submit\"**\n",
        "\n",
        "\n",
        "def prefill_form(src, fields: dict):\n",
        "  '''\n",
        "  src: the original src url to embed the form\n",
        "  fields: a dictionary of field:value pairs,\n",
        "  e.g. {\"pennkey\": my_pennkey, \"location\": my_location}\n",
        "  '''\n",
        "  prefill_fields = {}\n",
        "  for key in fields:\n",
        "      new_key = 'prefill_' + key\n",
        "      prefill_fields[new_key] = fields[key]\n",
        "  prefills = urllib.parse.urlencode(prefill_fields)\n",
        "  src = src + prefills\n",
        "  return src\n",
        "\n",
        "\n",
        "#autofill time if it is not present\n",
        "try: t0;\n",
        "except NameError: t0 = time.time()\n",
        "try: t1;\n",
        "except NameError: t1 = time.time()\n",
        "try: t2;\n",
        "except NameError: t2 = time.time()\n",
        "try: t3;\n",
        "except NameError: t3 = time.time()\n",
        "try: t4;\n",
        "except NameError: t4 = time.time()\n",
        "try: t5;\n",
        "except NameError: t5 = time.time()\n",
        "\n",
        "\n",
        "#autofill fields if they are not present\n",
        "#a missing pennkey and pod will result in an Airtable warning\n",
        "#which is easily fixed user-side.\n",
        "try: my_pennkey;\n",
        "except NameError: my_pennkey = \"\"\n",
        "try: my_pod;\n",
        "except NameError: my_pod = \"Select\"\n",
        "try: learning_from_previous_week;\n",
        "except NameError: learning_from_previous_week = \"\"\n",
        "try: value_estimate;\n",
        "except NameError: value_estimate = \"\"\n",
        "try: bad_actions;\n",
        "except NameError: bad_actions = \"\"\n",
        "try: grad_vars;\n",
        "except NameError: grad_vars = \"\"\n",
        "try: whitening_intuition;\n",
        "except NameError: whitening_intuition = \"\"\n",
        "try: obstacle;\n",
        "except NameError: obstacle = \"\"\n",
        "\n",
        "times = np.array([t1,t2,t3,t4,t5])-t0\n",
        "\n",
        "fields = {\"pennkey\": my_pennkey,\n",
        "          \"pod\": my_pod,\n",
        "          \"learning_from_previous_week\": learning_from_previous_week,\n",
        "          \"value_estimate\": value_estimate,\n",
        "          \"bad_actions\": bad_actions,\n",
        "          \"grad_vars\": grad_vars,\n",
        "          \"whitening_intuition\": whitening_intuition,\n",
        "          \"obstacle\": obstacle,\n",
        "          \"cumulative_times\": times}\n",
        "\n",
        "src = \"https://airtable.com/embed/shrflf6o9fqfsZhO9?\"\n",
        "\n",
        "\n",
        "# now instead of the original source url, we do: src = prefill_form(src, fields)\n",
        "display.display(IFrame(src = prefill_form(src, fields), width = 800, height = 400));"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qvyM4wuUfYCK"
      },
      "source": [
        "## Feedback\n",
        "How could this session have been better? How happy are you in your group? How do you feel right now?\n",
        "\n",
        "Feel free to use the embeded form below or use this link:\n",
        "<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://airtable.com/shrNSJ5ECXhNhsYss\">https://airtable.com/shrNSJ5ECXhNhsYss</a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RZMp_nI4fawU"
      },
      "source": [
        "display.display(IFrame(src=\"https://airtable.com/embed/shrNSJ5ECXhNhsYss?backgroundColor=red\", width = 800, height = 400))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qMxIV7tmq5fs"
      },
      "source": [
        "---\n",
        "# Appendix"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TqZbYv7kNJNE"
      },
      "source": [
        "## A.1 Derivation of Policy Gradient\n",
        "\n",
        "A typical reinforcement learning problem can be formulated as finding the optimial policy $\\pi^*$ that maximizes the expected reward. Consider a class of parametrized policy $\\{\\pi_{\\theta}\\}$. (If we use neural nets to approximate the policy, then $\\theta$ are the weights.) We can formalize the problem as\n",
        "\\begin{align*}\n",
        "\\theta^* &= \\underset{\\theta}{\\operatorname{argmax}} J(\\theta) \\\\\n",
        "J(\\theta) &= E_{\\tau \\sim p_{\\theta}(\\tau)}\\left[r(\\tau)\\right]  \n",
        "= \\int p_{\\theta}(\\tau) r(\\tau)d\\tau\n",
        "\\end{align*}\n",
        "<!-- \\begin{align*}\n",
        "&= \\int \\underbrace{p(s_1)\\prod_{t=1}^{T}\\pi_{\\theta}(a_t|s_t)p(s_{t+1}|s_{t}, a_{t})}_{p_{\\theta}(\\tau)} \\underbrace{\\sum_{t=1}^{T}r(s_t, a_t)}_{r(\\tau)}\n",
        "\\end{align*} -->\n",
        "\n",
        "In the objective, $\\tau$ denotes a trajectory, i.e. a sequence of state-action pairs $\\tau = (s_1, a_1), \\dots (s_T, a_T)$，where $T$ is commonly referred to as the horizon. $r(\\tau)$ is the reward of $\\tau$. $p_{\\theta}(\\tau)$ is the likelihood of $\\tau$ under the policy $\\pi_{\\theta}$.\n",
        "\\begin{align*}\n",
        "r(\\tau) &= \\sum_{t=1}^{T}r(s_t, a_t) \\\\\n",
        "p_{\\theta}(\\tau) &= p(s_1)\\prod_{t=1}^{T}\\pi_{\\theta}(a_t|s_t)p(s_{t+1}|s_{t}, a_{t})\n",
        "\\end{align*} \n",
        "\n",
        "To optimize the objective, we can try directly differentiating $J(\\theta)$ with respect to $\\theta$ and apply gradient ascent. \n",
        "\\begin{align*}\n",
        "\\nabla_{\\theta} J(\\theta) = \\nabla_{\\theta} \\int p_{\\theta}(\\tau) r(\\tau)d\\tau = \\int \\nabla_{\\theta}(p_{\\theta}(\\tau))r(\\tau)\n",
        "\\end{align*}\n",
        "However, the gradient is intractable because $p_{\\theta}(\\tau)$ involves the transition probability $p(s_{t+1}|s_t, a_t)$ which is unknown to us. We can instead apply the following convenient identity to get rid of the transition probability:\n",
        "\\begin{align*}\n",
        "\\nabla_{\\theta}(p_{\\theta}(\\tau)) = p_{\\theta}(\\tau)\\frac{\\nabla_{\\theta}(p_{\\theta}(\\tau))}{p_{\\theta}(\\tau)} = p_{\\theta}(\\tau)\\nabla_{\\theta}\\log p_{\\theta}(\\tau)\n",
        "\\end{align*}\n",
        "Let's plug it back into $\\nabla_{\\theta}J(\\theta)$ and see what happens. \n",
        "\n",
        "\\begin{align*}\n",
        "\\nabla_{\\theta}J(\\theta) &= \\int \\nabla_{\\theta}(p_{\\theta}(\\tau))r(\\tau) \\\\\n",
        "&= \\int p_{\\theta}(\\tau)\\nabla_{\\theta}\\log p_{\\theta}(\\tau)r(\\tau) \\\\\n",
        "&= E_{\\tau \\sim p_{\\theta}(\\tau)} \\left[\\nabla_{\\theta}\\log p_{\\theta}(\\tau)r(\\tau)\\right]\n",
        "\\end{align*}\n",
        "\n",
        "Expanding the log gradient term, we get \n",
        "\\begin{align*}\n",
        "    \\nabla_{\\theta}\\log p_{\\theta}(\\tau) &= \\nabla_{\\theta}\\log \\left(p(s_1)\\prod_{t=1}^{T}\\pi_{\\theta}(a_t|s_t)p(s_{t+1}|s_{t}, a_{t})\\right) \\\\\n",
        "    &=\\underbrace{\\nabla_{\\theta}\\log p(s_1)}_{0} + \\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t) + \\underbrace{\\nabla_{\\theta}\\log p(s_{t+1}|s_{t}, a_{t})}_{0} \\\\\n",
        "    &= \\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)\n",
        "\\end{align*}\n",
        "Viola! Using the log gradient identity, we managed to get rid of the transition probabilities! Now the policy gradient becomes \n",
        "\\begin{align*}\n",
        "\\nabla_{\\theta}J(\\theta) &= E_{\\tau \\sim p_{\\theta}(\\tau)} \\left[\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)\\sum_{t=1}^{T}r(s_t, a_t)\\right] \\\\\n",
        "&\\approx \\frac{1}{N} \\sum_{i=1}^{N} \\left(\\sum_{t=1}^{T}\\nabla_{\\theta}\\log\\pi_{\\theta}(a_t|s_t)\\sum_{t=1}^{T}r(s_t, a_t)\\right)\n",
        "\\end{align*}\n",
        "\n",
        "The expectation can be evaluated via Monte-Carlo sampling, which is just a fancy way to say \"take the average of many samples.\" How do we evaluate the $\\log\\pi_{\\theta}(a_t|s_t)$ term? If the action space is discrete, then we can interpret the output of the neural network as the logits of actions and obtain the probabilities by taking the softmax. If the action space is continuous, then we can parametrize the policy using a standard probability distribution (e.g. Gaussian) which will naturally give us the probability of an action.\n",
        "\n",
        "To get some intuition, let's look more closely at the policy gradient in it's simplified form. $$\\nabla_{\\theta}J(\\theta) \\approx \\frac{1}{N}\\sum_{i=1}^N\\nabla_{\\theta}\\log\\pi_{\\theta}(\\tau_i)r(\\tau_i)$$ Note that we are essentially weighting the gradient of each trajectory proportionally to its reward. If a trajectory $\\tau$ has a high reward, then we increase the likelihood of the actions seen in $\\tau$. Otherwise we decrease the likelihood of those actions. It may seem imprudent to say that all actions in a high-reward trajectory are good, but in expectation this update converges to the optimal policy."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5NbtzGC0cegD"
      },
      "source": [
        "## A.2 Actor-critic from a TD perspective. \n",
        "\n",
        "Actor-critic methods are TD methods that have a separate memory structure to explicitly represent the policy independent of the value function. The policy structure is known as the actor, because it is used to select actions, and the estimated value function is known as the critic, because it critiques the actions made by the actor. Learning is always on-policy: the critic must learn about and critique whatever policy is currently followed by the actor. Typically, the critic is a state-value function. After each action selection, the critic evaluates the new state to determine whether things have gone better or worse than expected. That evaluation is the TD error:\n",
        "\n",
        "$$\n",
        "\\delta_{t}=r_{t+1}+\\gamma V\\left(s_{t+1}\\right)-V\\left(s_{t}\\right)\n",
        "$$\n",
        "\n",
        "where $V(s)$ is the current value function implemented by the critic. This TD error can be used to evaluate the action $a_t$ taken in state $s_t$. If the TD error is positive, it suggests that the tendency to select $a_t$ should be strengthened for the future, whereas if the TD error is negative, it suggests the tendency should be weakened. \n",
        "\n",
        "Advantage function is used to stabilize learning. It is defined as:\n",
        "$$\n",
        "A(s,a) = Q(s,a) - V(s)\n",
        "$$\n",
        "where $Q(s,a)$ is the Q-value for action a in state $s$, and $V(s)$ is the value of that state. The function tells us the improvement compared to the average value of the action taken at that state. If $A(s,a) > 0$, our gradient is pushed in that direction, whereas if $A(s,a) < 0$, our gradient is pushed in the opposite direction. The TD error is a good estimator of the advantage function: \n",
        "\n",
        "$$ \n",
        "A(s,a) = Q(s,a) - V(s) \\approx r + \\gamma V(s') - V(s) \n",
        "$$"
      ]
    }
  ]
}